package com.accesslimit.demo;

import java.time.LocalTime;
import java.util.HashSet;
import java.util.Set;
import java.util.TreeMap;

/**
 * 滑动日志算法
 *
 * @author sam
 * @description 代码比滑动窗口简单，概念差不多，都是动态地判断时间范围，测试结果比滑动窗口更准。滑动日志算法缺点：占用内存更多
 */
public class RateLimiterSildingLog {

    // 阈值
    private Integer qps = 2;

    // 记录当前请求时间戳和数量
    private TreeMap<Long, Long> treeMap = new TreeMap<>();

    // 清理请求日志间隔时间 60s
    private long clearTime = 60 * 1000;

    public RateLimiterSildingLog(Integer qps) {
        this.qps = qps;
    }

    private synchronized boolean tryAcquire() {
        long now = System.currentTimeMillis();
        // 请求旧数据
        if (!treeMap.isEmpty() && (treeMap.firstKey() - now) > clearTime) {
            Set<Long> keySet = new HashSet<>(treeMap.subMap(0L, now - 1000).keySet());
            for (Long key : keySet) {
                treeMap.remove(key);
            }
        }
        // 计算当前请求次数
        int sum = 0;
        for (Long value : treeMap.subMap(now - 1000, now).values()) {
            sum += value;
        }
        // 超过qps限制返回
        if (sum >= qps) {
            return false;
        }
        // 记录本次日志
        if (treeMap.containsKey(now)) {
            // +1操作
            treeMap.compute(now, (k, v) -> v + 1);
        } else {
            // 初次放入操作
            treeMap.put(now, 1L);
        }
        return true;
    }

    public static void main(String[] args) throws InterruptedException {
        RateLimiterSildingLog rateLimiterSildingLog = new RateLimiterSildingLog(3);
        for (int i = 0; i < 50; i++) {
            Thread.sleep(250);
            LocalTime now = LocalTime.now();
            if (rateLimiterSildingLog.tryAcquire()) {
                System.out.println(now + ": " + "该请求已通过");
            } else {
                System.out.println(now + ": " + "该请求已被限流！");
            }
        }
    }

}
